﻿#pragma once

#include <optixu/optixu_math_namespace.h>
#include <optix.h>
#include <optixu/optixu_math_namespace.h>
#include "helpers.h"
#include "path_tracer.h"
#include "materials.h"
#include "random.h"
#include "sunsky.h"
#include "materials.h"
#include "lights_sampling.h"
#include "bsdf_sampling.h"
#include "Volume.h"

using namespace optix;

rtDeclareVariable(float, scene_epsilon, , );
rtDeclareVariable(rtObject, top_object, , );
rtDeclareVariable(optix::Ray, ray, rtCurrentRay, );
rtDeclareVariable(float, t_hit, rtIntersectionDistance, );
rtDeclareVariable(unsigned int, pathtrace_ray_type, , );
rtDeclareVariable(unsigned int, pathtrace_shadow_ray_type, , );
rtDeclareVariable(unsigned int, rr_begin_depth, , );
rtDeclareVariable(unsigned int, max_depth, , );
rtDeclareVariable(unsigned int, use_envmap, , ) = 0;
rtBuffer<TriangleLight>     lights;

//geometry attributes
rtDeclareVariable(float3, geometric_normal, attribute geometric_normal, );
rtDeclareVariable(float3, shading_normal, attribute shading_normal, );
rtDeclareVariable(float3, texcoord, attribute texcoord, );
rtDeclareVariable(float3, emission_color, , );
rtDeclareVariable(float3, diffuse_color, , );
rtDeclareVariable(float3, bg_color, , );
rtDeclareVariable(int, use_texture, , ) = 0;
rtDeclareVariable(float, extinction, , ) = -0.06f;

rtDeclareVariable(float, phong_exp, , );
rtDeclareVariable(float3, directional_light, , );
rtDeclareVariable(float3, directional_light_col, , );
rtDeclareVariable(float, solid_angle, , ) = 6.0e-5f * 100.0f;

rtDeclareVariable(PerRayData_pathtrace, current_prd, rtPayload, );

rtDeclareVariable(PerRayData_pathtrace_shadow, current_prd_shadow, rtPayload, );
rtDeclareVariable(float3, shadow_attenuation, , );
rtDeclareVariable(float, importance_cutoff, , ) = 0.0025f;

//Sample Infinite Light
rtTextureSampler<float4, 2> envmap;
rtTextureSampler<float4, 2> diffuse_map;
rtTextureSampler<float4, 2> envmap_lt;

__device__ __inline__ optix::float4 rect_envmap_le(optix::float3& dir)
{
	float3 direction = -dir;
	float theta = atan2f(direction.x, direction.y);
	float phi = M_PIf * 0.5f - acosf(direction.z);
	float u = (theta + M_PIf) * (0.5f * M_1_PIf);
	float v = 0.5f * (1.0f + sin(phi));

	return tex2D(envmap_lt, u, v);
}


__device__ __inline__ void sample_lights(const optix::float3 &sample, const optix::float3 hitpoint, const optix::float3 ffnormal, optix::float3 &L, optix::float3 &Le, float &ndl, float &Ldist, float&pdf)
{
	using namespace optix;

	if (lights.size() == 0)
	{
		if (use_envmap > 0)
		{
			sample_envmap(make_float2(sample), ffnormal, L, ndl);
			Ldist = 100.0f;
			Le = make_float3(rect_envmap_le(L));
			pdf = ndl;
			return;
		}
		sample_sunsky(solid_angle, make_float2(sample), ffnormal, L, ndl);
		pdf = ndl;
		Le = directional_light_col;
		return;
	}

	const unsigned int num_lights = lights.size() - 1;
	{
		const float z1 = sample.x;
		const float z2 = sample.y;
		TriangleLight light = lights[sample.z*(num_lights)];
		float alpha = 1.0f - sqrt(fold(z1));
		float beta = (1.0f - fold(z2)) * sqrt(fold(z1));
		float gamma = fold(z2) * sqrt(fold(z1));
		float3 light_pos = light.v3 * gamma + light.v1 * alpha + light.v2 * beta;
		 Ldist = length(light_pos - hitpoint);
		L = normalize(light_pos - hitpoint);
		ndl = dot(ffnormal, L);
		float LnDl = dot(light.normal, L);
		float A = length(cross(light.v1 - light.v2, light.v1 - light.v3));
		pdf = ndl * LnDl / (M_PIf*Ldist*Ldist) * A;

		Le = make_float3(1, 1, 1)*num_lights * light.emission;

		if (LnDl < 0.f) {
			ndl = -1;
		}
	}
}
